import base64
import io
import os
import random
import string
import time
from enum import Enum

from PIL import Image
from captcha.image import ImageCaptcha

from utils.helper import uuid
from utils.paths import DefaultCaptchaPath


class CaptchaOutputType(Enum):
    file = "file"
    base64 = "base64"

captcha = ImageCaptcha()

def generate_captcha(path=None, output: CaptchaOutputType = CaptchaOutputType.file) -> (str, str):
    # 删除5分钟前的验证码图片
    current_time = time.time()
    for root, dirs, files in os.walk(DefaultCaptchaPath):
        for file in files:
            file_path = os.path.join(root, file)
            file_modify_time = os.path.getmtime(file_path)
            time_difference = current_time - file_modify_time
            if time_difference > 300:
                os.remove(file_path)

    if not path:
        path = DefaultCaptchaPath

    file_name = uuid() + ".png"
    file_path = str(os.path.join(path, file_name))

    captcha_text = ''.join(random.choices(string.ascii_letters + string.digits, k=4))
    captcha.write(captcha_text, file_path)

    if output == CaptchaOutputType.file:
        return file_path
    elif output == CaptchaOutputType.base64:
        with Image.open(file_path) as img:
            buffer = io.BytesIO()
            img.save(buffer, format="PNG")
            img_bytes = buffer.getvalue()

        base64_str = base64.b64encode(img_bytes).decode("utf-8")

        return {"text": captcha_text, "base64": "data:image/png;base64," + base64_str}


if __name__ == "__main__":
    data = generate_captcha(output=CaptchaOutputType.base64)
    print(data)
